import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout, BatchNormalization
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import time

class LSTMPredictor:
    """LSTM时序预测模型，用于价格波动预警"""
    def __init__(self, config=None):
        self.config = config or {
            'sequence_length': 60,  # 序列长度
            'batch_size': 32,
            'epochs': 100,
            'learning_rate': 0.001,
            'hidden_units': [64, 32],
            'dropout_rate': 0.2
        }
        
        self.models = {}
        self.scalers = {}
        self.prediction_history = {}
        self.cycle_times = {'15s': 15, '1m': 60, '5m': 300}  # 不同周期的秒数
        
    def build_model(self, input_shape):
        """构建LSTM模型"""
        model = Sequential()
        
        # 第一层LSTM
        model.add(LSTM(
            units=self.config['hidden_units'][0],
            return_sequences=True,
            input_shape=input_shape
        ))
        model.add(BatchNormalization())
        model.add(Dropout(self.config['dropout_rate']))
        
        # 第二层LSTM
        model.add(LSTM(
            units=self.config['hidden_units'][1],
            return_sequences=False
        ))
        model.add(BatchNormalization())
        model.add(Dropout(self.config['dropout_rate']))
        
        # 输出层
        model.add(Dense(units=1))
        
        # 编译模型
        optimizer = tf.keras.optimizers.Adam(learning_rate=self.config['learning_rate'])
        model.compile(optimizer=optimizer, loss='mse')
        
        return model
    
    def prepare_data(self, data, cycle='1m'):
        """准备LSTM输入数据"""
        # 检查数据是否为空
        if data is None or len(data) < self.config['sequence_length'] + 1:
            return None, None
        
        # 根据周期对数据进行采样
        cycle_time = self.cycle_times.get(cycle, 60)
        sampled_data = data.iloc[::int(cycle_time/self._get_data_frequency(data))]
        
        # 提取价格列
        prices = sampled_data['price'].values.reshape(-1, 1)
        
        # 初始化或使用已有的scaler
        symbol = data['symbol'].iloc[0] if 'symbol' in data.columns else 'default'
        key = f"{symbol}_{cycle}"
        
        if key not in self.scalers:
            self.scalers[key] = MinMaxScaler(feature_range=(0, 1))
            scaled_prices = self.scalers[key].fit_transform(prices)
        else:
            scaled_prices = self.scalers[key].transform(prices)
        
        # 创建序列数据
        X, y = [], []
        seq_len = self.config['sequence_length']
        
        for i in range(seq_len, len(scaled_prices)):
            X.append(scaled_prices[i-seq_len:i, 0])
            y.append(scaled_prices[i, 0])
        
        # 转换为numpy数组并调整形状
        X = np.array(X)
        y = np.array(y)
        
        # LSTM需要的输入形状是 [samples, time steps, features]
        X = np.reshape(X, (X.shape[0], X.shape[1], 1))
        
        return X, y
    
    def _get_data_frequency(self, data):
        """估算数据的频率（秒）"""
        if 'timestamp' not in data.columns:
            return 1  # 默认假设1秒频率
        
        # 计算时间戳的差值
        timestamps = data['timestamp'].values
        if len(timestamps) < 2:
            return 1
        
        # 转换为秒
        if timestamps[0] > 1e9:  # 毫秒时间戳
            timestamps = timestamps / 1000
        
        # 计算平均间隔
        intervals = np.diff(timestamps)
        avg_interval = np.mean(intervals)
        
        return max(1, int(avg_interval))  # 确保至少为1秒
    
    def train(self, data, symbol, cycle='1m'):
        """训练模型"""
        # 准备数据
        X, y = self.prepare_data(data, cycle)
        if X is None or y is None:
            print(f"Insufficient data to train model for {symbol} at {cycle} cycle")
            return False
        
        # 构建或获取模型
        key = f"{symbol}_{cycle}"
        if key not in self.models:
            self.models[key] = self.build_model((X.shape[1], 1))
        
        model = self.models[key]
        
        # 设置回调函数
        callbacks = [
            EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True),
            ModelCheckpoint(f'models/{key}_lstm_model.h5', monitor='val_loss', save_best_only=True)
        ]
        
        # 分割训练集和验证集
        train_size = int(len(X) * 0.8)
        X_train, X_val = X[:train_size], X[train_size:]
        y_train, y_val = y[:train_size], y[train_size:]
        
        # 训练模型
        history = model.fit(
            X_train, y_train,
            batch_size=self.config['batch_size'],
            epochs=self.config['epochs'],
            validation_data=(X_val, y_val),
            callbacks=callbacks,
            verbose=0
        )
        
        print(f"Model for {symbol} at {cycle} cycle trained successfully")
        return True
    
    def predict(self, data, symbol, cycle='1m'):
        """预测未来价格"""
        key = f"{symbol}_{cycle}"
        
        # 检查模型是否存在
        if key not in self.models or key not in self.scalers:
            print(f"Model not found for {symbol} at {cycle} cycle")
            return None
        
        model = self.models[key]
        scaler = self.scalers[key]
        
        # 准备最新数据用于预测
        seq_len = self.config['sequence_length']
        
        # 根据周期采样数据
        cycle_time = self.cycle_times.get(cycle, 60)
        freq = self._get_data_frequency(data)
        sampled_data = data.iloc[::int(cycle_time/freq)]
        
        if len(sampled_data) < seq_len:
            print(f"Insufficient data for prediction for {symbol} at {cycle} cycle")
            return None
        
        # 提取最新的序列数据
        latest_prices = sampled_data['price'].values[-seq_len:].reshape(-1, 1)
        scaled_prices = scaler.transform(latest_prices)
        
        # 调整形状用于模型输入
        X_pred = np.reshape(scaled_prices, (1, seq_len, 1))
        
        # 进行预测
        start_time = time.time()
        scaled_prediction = model.predict(X_pred, verbose=0)[0][0]
        predict_time = (time.time() - start_time) * 1000  # 毫秒
        
        # 反归一化预测结果
        prediction = scaler.inverse_transform([[scaled_prediction]])[0][0]
        
        # 记录预测历史
        if key not in self.prediction_history:
            self.prediction_history[key] = []
        
        current_time = int(time.time() * 1000)
        self.prediction_history[key].append({
            'timestamp': current_time,
            'predicted_price': prediction,
            'current_price': latest_prices[-1][0],
            'predict_time_ms': predict_time
        })
        
        # 限制历史记录大小
        if len(self.prediction_history[key]) > 1000:
            self.prediction_history[key] = self.prediction_history[key][-1000:]
        
        # 计算价格波动预警指标
        price_change = prediction - latest_prices[-1][0]
        percent_change = (price_change / latest_prices[-1][0]) * 100
        
        return {
            'symbol': symbol,
            'predicted_price': prediction,
            'current_price': latest_prices[-1][0],
            'price_change': price_change,
            'percent_change': percent_change,
            'cycle': cycle,
            'timestamp': current_time,
            'predict_latency_ms': predict_time
        }
    
    def load_model(self, symbol, cycle='1m', model_path=None):
        """加载预训练模型"""
        key = f"{symbol}_{cycle}"
        
        try:
            path = model_path or f'models/{key}_lstm_model.h5'
            self.models[key] = tf.keras.models.load_model(path)
            print(f"Model loaded successfully for {symbol} at {cycle} cycle")
            return True
        except Exception as e:
            print(f"Failed to load model for {symbol} at {cycle} cycle: {str(e)}")
            return False
    
    def get_prediction_accuracy(self, symbol, cycle='1m'):
        """获取预测准确率"""
        key = f"{symbol}_{cycle}"
        
        if key not in self.prediction_history or len(self.prediction_history[key]) < 10:
            return None
        
        # 计算预测准确率指标（简单的方向性准确率）
        history = self.prediction_history[key][-50:]  # 最近50次预测
        correct_direction = 0
        
        for i in range(1, len(history)):
            # 预测方向
            predicted_direction = 1 if history[i]['predicted_price'] > history[i-1]['current_price'] else -1
            # 实际方向
            actual_direction = 1 if history[i]['current_price'] > history[i-1]['current_price'] else -1
            
            if predicted_direction == actual_direction:
                correct_direction += 1
        
        accuracy = correct_direction / (len(history) - 1) if len(history) > 1 else 0
        
        return {
            'symbol': symbol,
            'cycle': cycle,
            'accuracy': accuracy,
            'samples': len(history) - 1
        }